from netCDF4 import Dataset
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from matplotlib.cm import get_cmap
#import cartopy.crs as crs
#from cartopy.feature import NaturalEarthFeature
import os
#from mpl_toolkits.basemap import Basemap
import geopy.distance

# 2, 6, 7,
################################################################################################
#functions
################################################################################################

# This function returns a list of all wrf files in the directory.
def list_files(Dir, ncfiles):
	for f in os.listdir(Dir):
		if f.endswith('CHRTOUT_DOMAIN1'):									#edit here
			ncfiles.append(Dir+'/'+f)
	return (ncfiles)

def mapper(lat, lon, name):
    label = '('+str(round(lat, 2))+', '+str(round(lon,2))+')'
    m = Basemap(llcrnrlon = ll_lon, llcrnrlat = ll_lat, urcrnrlon = ur_lon, urcrnrlat = ur_lat,
                rsphere=(6378137.00,6356752.3142),\
                    resolution='h',area_thresh=0.1,projection='lcc',\
                        lat_1=30,lat_2=60,lat_0=40,lon_0=-97,lon_1 = -97)

    try:
        #m.arcgisimage(service='World_Shaded_Relief')
        m.drawrivers(color='dodgerblue',linewidth= 2.0,zorder=5)
        lons, lats = m(lon, lat)
        m.scatter(lons, lats, marker = 'o', color='r', zorder=5)
   
        #m.drawcoastlines(linewidth=0.5)
        m.drawstates(linewidth = 0.2)
        #m.drawcounties(linewidth=0.1)
        m.shadedrelief()
        #m.drawlsmask()
        #m.etopo()
        #m.fillcontinents(color='coral',lake_color='aqua')
        #m.drawmapboundary(fill_color='aqua')
    except:
        pass

    # draw parallels and meridians.
    parallels = np.arange(-80.,81.,2)
    meridians = np.arange(-180.,181.,2)
    # labels = [left,right,top,bottom]
    m.drawparallels(parallels,labels=[False,True,True,False])
    m.drawmeridians(meridians,labels=[True,False,False,True])
    #plt.figure()
    plt.title(name)
    plt.text(lons, lats, label,fontsize=12,fontweight='bold', ha='left',va='bottom')
    plt.show()
    #plt.savefig(name+'.png')

##############################################
# Getting the netcdfs
##############################################
Dir = '/home/mkhondak/project/momen/murad/WRF-Hydro_Offline/WRF-Hydro/Hurricane/harvey/larger_houston/cuahsi_modified_houston_changed/WRF_WRF-Hydro/special_4/Feb_24_def_rd_redo//8km/feb_2023/pbl_1_25_0'		#edit here
ncfiles = []

# input NetCDF file

# Use the list_files function to list all the wrf files in the directory.
ncfiles = list_files(Dir, ncfiles)

# Sort the ncfiles 
ncfiles = sorted(ncfiles)
#print(ncfiles)

temp = Dataset(ncfiles[60])

################################################################################################
# Input
################################################################################################

timeseries_flag = 1

gauge_file = 'gauges.csv'

gauge_dataset = pd.read_csv(gauge_file,sep='\t')									#edit here
list_name = gauge_dataset.iloc[:, 2].values
list_lat = gauge_dataset.iloc[:, 4].values
list_lon = gauge_dataset.iloc[:, 5].values


# Boundary Box

ll_lon = -99
ll_lat = 27.3 
ur_lon = -92.9 
ur_lat = 33.4

# list_lat = [30.34971, 30.031986]
# list_lon = [-94.08446, -95.26311]

################################################################################################
# Main Code
################################################################################################

##############################################
# getting index  for forecast point
##############################################


# showing the variables stored in this file
#print(temp.variables.keys())
#print(temp.variables['feature_id'])
#print(temp.variables['latitude'].shape)


# Get the values in respective variables
lats = temp.variables['latitude'][:]
lons = temp.variables['longitude'][:]
fid = temp.variables['feature_id'][:]
order = temp.variables['order'][:]
#head = temp.variables['Head'][:]
streamflow = temp.variables['streamflow'][:]
time = temp.model_output_valid_time
no_fid = len(fid)



for i in range(len(list_lat)):                 # 1-10 doing
    file = str(i+1)
    file_name = file+'_YSU_default_diffusion_streamflow.txt'										#edit here
    in_lat = list_lat[i] 
    in_long =  list_lon[i]
    in_name = list_name[i]
# Getting the closest location from the given location
    index_frxst = ((np.abs(lons - in_long))+(np.abs(lats - in_lat))).argmin()
    #distance = ((lons[index_frxst] - in_long)**2 +  (lats[index_frxst] - in_lat)**2)**0.5
    distance = geopy.distance.geodesic((in_lat, in_long), (lats[index_frxst],lons[index_frxst])).km

    print('the input coordinate (',in_lat, in_long, ') \n the nearest coordinate is (', lats[index_frxst], lons[index_frxst], ')\n the distance is', distance, '\n the index is' , index_frxst)

    #print('check the value in frxst_pts_out.txt to verify codes ok!\n', index_frxst, time, lats[index_frxst], lons[index_frxst], streamflow[index_frxst], '\n')



##################################################
# timeseries for that forecast point
##################################################
    if timeseries_flag ==1:
        with open(file_name, 'w') as f:
            #f.write('the input coordinate ('+str(in_lat)+' '+str(in_long)+',    the nearest coordinate is ('+' '+str(lats[index_frxst])+' '+str(lons[index_frxst])+'),    the distance is'+str(distance)+',   the index is'+str(index_frxst)+'\n')
            for item in range(len(ncfiles)):
                data = Dataset(ncfiles[item])
                lats_d = data.variables['latitude'][:]
                lons_d = data.variables['longitude'][:]
                fid_d = data.variables['feature_id'][:]
                # head_d = data.variables['Head'][:]
                streamflow_d = data.variables['streamflow'][:]
                q_lateral_d = data.variables['q_lateral'][:]
                velocity_d = data.variables['velocity'][:]
                elevation_d = data.variables['elevation'][:]
                order_d = data.variables['order'][:]
                time_d = data.model_output_valid_time
                print(time_d, lats_d[index_frxst], lons_d[index_frxst], order_d[index_frxst], streamflow_d[index_frxst])
                f.write(str(time_d)+' '+str(lats_d[index_frxst])+' '+str(lons_d[index_frxst])+' '+str(order_d[index_frxst])+' '+str(streamflow_d[index_frxst])+'\n')
            f.close()
    

